#!/usr/bin/env python3
# Author: Armit
# Create Time: 2023/01/11 

# 反正 cd-hit 可以对 pdb 文件进行聚类

from re import compile as Regex
from pathlib import Path
from typing import List, Tuple

from data import RDATA_PATH, DATA_PATH

TITLE_REGEX = Regex('>Cluster (\d+)')
ITEM_REGEX1 = Regex('\d+	\d+aa, >(\w{4})... \*')
ITEM_REGEX2 = Regex('\d+	\d+aa, >(\w{4})... at ([\d\.]+)%')


def preprocess(in_fn, out_fn):
  print(f'>> process {in_fn!r}')
  with open(Path(RDATA_PATH) / in_fn, 'r', encoding='utf-8') as fh:
    lines = fh.read().strip().split('\n')
  
  clusters:List[List[Tuple[float, str]]] = []
  cluster, sn = [], None
  for line in lines:
    m = TITLE_REGEX.match(line)
    if m:
      if cluster: clusters.append(cluster)
      cluster = []
    else:
      m2 = ITEM_REGEX1.match(line)
      if m2:
        pdb_id = m2.groups()[0]
        sim_sc = 100.0
      else:
        m3 = ITEM_REGEX2.match(line)
        id_sc = m3.groups()
        pdb_id = id_sc[0]
        sim_sc = float(id_sc[1])
      cluster.append((sim_sc, pdb_id))
  clusters.append(cluster)
  del cluster, sn

  print(f'>> found {sum(len(c) for c in clusters)} pdb record')
  n_clusters = len(clusters)
  print(f'>> found {n_clusters} clusters')

  clusters = sorted(clusters, key=lambda x: len(x), reverse=True)
  print(f'      max cluster size: {len(clusters[0])}')
  print(f'      min cluster size: {len(clusters[-1])}')

  with open(Path(DATA_PATH) / out_fn, 'w', encoding='utf-8') as fh:
    for cluster in clusters:
      for sc, id in sorted(cluster, reverse=True):
        fh.write(f'({id}, {sc}) ')
      fh.write('\n')


if __name__ == '__main__':
  preprocess('Single.fas.clstr', 'cdhit_clusters_s.txt')
  preprocess('Multi.fas.clstr',  'cdhit_clusters_m.txt')
